[Perf] Early return in KVCacheManager.allocate_slots #29206
Conversation
Signed-off-by: Jialin Ouyang <[email protected]>
Jialin
requested review from
ApostaC,
WoosukKwon,
alexm-redhat,
njhill,
robertgshaw2-redhat and
ywang96
as code owners
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if ( | ||
| num_blocks_to_allocate == 0 | ||
| and new_computed_block_list is self.empty_kv_cache_blocks.blocks | ||
| ): | ||
| # Early return as no new blocks needed to be allocated | ||
| return self.empty_kv_cache_blocks |
There was a problem hiding this comment.
allocate_slots now returns early when num_blocks_to_allocate==0 and no prefix-cache hits, so coordinator.remove_skipped_blocks() is never called on those steps. For sliding-window/ChunkedLocal attention this call is what frees blocks that have fallen outside the window; skipping it leaves those blocks held until the request hits a block boundary or finishes, inflating block_pool usage and causing unnecessary evictions or scheduling failures for long prompts that keep generating within an existing block.
Useful? React with 👍 / 👎.
|
|
||
| if ( | ||
| num_blocks_to_allocate == 0 | ||
| and new_computed_block_list is self.empty_kv_cache_blocks.blocks |
There was a problem hiding this comment.
| and new_computed_block_list is self.empty_kv_cache_blocks.blocks | |
| and new_computed_blocks is None |
prefer this slightly.
And do we need to add num_computed_tokens > request.num_prompt_tokens to make remove_skipped_blocks to be called in the first decode step? This can help to free the prefill tokens used by the last prefill step but is outside the sliding window of the first decode step. Would be grateful if you can try gpt-oss and gemma3, two models with small sliding window size.
There was a problem hiding this comment.
And can you add a comment for this may delay remove_skipped_blocks and cache_blocks and give some analysis about it is fine.
There was a problem hiding this comment.
Hmmm, great call. I didn't notice the cache_blocks logic. Let me think about the early return logic more :/
Signed-off-by: Jialin Ouyang <[email protected]>
2 participants
Purpose
As the title, early return in KVCacheManager.allocate_slots to speed up scheduling, as most of the requests only allocate a block every 16 steps.
Test Plan && Test Result
With the change, we could observed allocate slots cost is cut by half from 17.24% to 8.24% over the total runtime (with async scheduling enabled)
Total runtime distribution for allocate_slots per engine step
Public benchmarks
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.